package com.central.gateway.filter;


import com.central.common.utils.IdGenerator;
import com.central.gateway.common.Constants;
import com.central.gateway.common.LogVo;
import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.NettyWriteResponseFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.factory.GatewayFilterFactory;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.io.UnsupportedEncodingException;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.util.Objects;

import static org.springframework.cloud.gateway.support.GatewayToStringStyler.filterToStringCreator;

/**
 * 获取参数请求
 * @author zouzhihui
 */
@Component
@Slf4j
public class ParameterGatewayFilterFactory extends AbstractGatewayFilterFactory<ParameterGatewayFilterFactory.Config> {


    private static final String POST = "POST";
    private static final String PUT = "";
    private static final String GET = "GET";
    private static final String UTF_8 = "utf-8";
    private static final String START_TIME = "startTime";
    private static final String QUERY_PARAMETER = "queryParameter";
    private static final String BODY_PARAMETER = "bodyParameter";
    private static final String TRACE_ID = "x-trace-id";

    public ParameterGatewayFilterFactory() {
        super(ParameterGatewayFilterFactory.Config.class);
    }

    @Override
    public GatewayFilter apply(Config config) {
        ParameterGatewayFilter parameterGatewayFilter = new ParameterGatewayFilter(config);
        parameterGatewayFilter.setGatewayFilterFactory(this);
        return parameterGatewayFilter;
    }

    public static class Config {

        private Boolean enableLog;

        public Boolean getEnableLog() {
            return enableLog;
        }

        public void setEnableLog(Boolean enableLog) {
            this.enableLog = enableLog;
        }
    }

    public class ParameterGatewayFilter implements GatewayFilter, Ordered {

        private final Config config;

        private GatewayFilterFactory<Config> gatewayFilterFactory;

        public ParameterGatewayFilter(Config config) {
            this.config = config;
        }

        @Override
        public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
            ServerHttpRequest serverHttpRequest = exchange.getRequest();
            ServerHttpRequest build = exchange.getRequest().mutate().header(TRACE_ID, IdGenerator.getIdStr()).build();
            if (Boolean.TRUE.equals(config.getEnableLog())) {
                String method = serverHttpRequest.getMethodValue();
                if (POST.equals(method)){
                    return DataBufferUtils.join(exchange.getRequest().getBody()).flatMap(dataBuffer -> {
                        byte[] bytes = new byte[dataBuffer.readableByteCount()];
                        dataBuffer.read(bytes);
                        String bodyString = "";
                        try {
                            bodyString = new String(bytes, UTF_8);
                        } catch (UnsupportedEncodingException e) {
                            e.printStackTrace();
                        }
                        DataBufferUtils.release(dataBuffer);
                        Flux<DataBuffer> cachedFlux = Flux.defer(() -> {
                            DataBuffer buffer = exchange.getResponse().bufferFactory().wrap(bytes);
                            return Mono.just(buffer);
                        });
                        ServerHttpRequest mutatedRequest = new ServerHttpRequestDecorator(build) {
                            @Override
                            public Flux<DataBuffer> getBody() {
                                return cachedFlux;
                            }
                        };
                        LogVo logVo = new LogVo(exchange.getRequest().getHeaders().getFirst(TRACE_ID), exchange.getRequest().getMethod().name(), exchange.getRequest().getURI().toString(),
                                exchange.getRequest().getQueryParams().toString(), bodyString, System.currentTimeMillis());
                        exchange.getAttributes().put(Constants.LOG_VO_ATTR, logVo);
                        return chain.filter(exchange.mutate().request(mutatedRequest).response(decorateResponse(exchange)).build());
                    });
                } else {
                    LogVo logVo = new LogVo(exchange.getRequest().getHeaders().getFirst(TRACE_ID), exchange.getRequest().getMethod().name(), exchange.getRequest().getURI().toString(),
                            exchange.getRequest().getQueryParams().toString(), System.currentTimeMillis());
                    exchange.getAttributes().put(Constants.LOG_VO_ATTR, logVo);
                }
            }
            return chain.filter(exchange.mutate().response(decorateResponse(exchange)).build());
        }

        @Override
        public int getOrder() {
            return NettyWriteResponseFilter.WRITE_RESPONSE_FILTER_ORDER - 1;
        }

        public void setGatewayFilterFactory(GatewayFilterFactory<Config> gatewayFilterFactory) {
            this.gatewayFilterFactory = gatewayFilterFactory;
        }


        private ServerHttpResponse decorateResponse(ServerWebExchange exchange) {
            return new ServerHttpResponseDecorator(exchange.getResponse()) {
                @Override
                public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                    exchange.getResponse().getHeaders().add(TRACE_ID, exchange.getRequest().getHeaders().getFirst(TRACE_ID));
                    if (body instanceof Flux) {
                        Flux<? extends DataBuffer> fluxBody = Flux.from(body);
                        return super.writeWith(fluxBody.buffer().map(dataBuffer -> {
                            DataBuffer join = exchange.getResponse().bufferFactory().join(dataBuffer);
                            byte[] content = new byte[join.readableByteCount()];
                            join.read(content);
                            DataBufferUtils.release(join);
                            String responseResult = new String(content, getMediaTypeCharset(exchange.getResponse().getHeaders().getContentType()));
                            //printPostLog(exchange, responseResult);
                            Object attribute = exchange.getAttribute(Constants.LOG_VO_ATTR);
                            if (Objects.nonNull(attribute)) {
                                LogVo logVo = (LogVo) attribute;
                                logVo.setEndTime(System.currentTimeMillis());
                                logVo.setResponseResult(responseResult);
                                log.info(logVo.toString());
                                exchange.getAttributes().remove(Constants.LOG_VO_ATTR);
                            }
                            return exchange.getResponse().bufferFactory().wrap(content);
                        }));
                    }
                    return super.writeWith(body);
                }

                @Override
                public Mono<Void> writeAndFlushWith(Publisher<? extends Publisher<? extends DataBuffer>> body) {
                    return writeWith(Flux.from(body).flatMapSequential(p -> p));
                }
            };
        }

        private Charset getMediaTypeCharset(MediaType mediaType) {
            if (Objects.nonNull(mediaType) && Objects.nonNull(mediaType.getCharset())) {
                return mediaType.getCharset();
            } else {
                return StandardCharsets.UTF_8;
            }
        }

        @Override
        public String toString() {
            Object obj = (this.gatewayFilterFactory != null) ? this.gatewayFilterFactory : this;
            return filterToStringCreator(obj)
                    .append("enableLog", config.getEnableLog()).toString();
        }
    }
}
